# -*- coding: utf-8 -*-
import time
import numpy as np
import h5py
import matplotlib.pyplot as plt
import scipy
from PIL import Image
from scipy import ndimage
from dnn_app_utils_v2 import *
# from C1_4_homework import *

plt.rcParams['figure.figsize'] = (5.0, 4.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

np.random.seed(1)

train_x_orig, train_y, test_x_orig, test_y, classes = load_data()
# index = 10
# plt.imshow(train_x_orig[index])
plt.show()

m_train = train_x_orig.shape[0] # m_train = 209
num_px = train_x_orig.shape[1]
m_test = test_x_orig.shape[0] # m_test =

# print ("Number of training examples: " + str(m_train))
# print ("Number of testing examples: " + str(m_test))
# print ("Each image is of size: (" + str(num_px) + ", " + str(num_px) + ", 3)")
# print ("train_x_orig shape: " + str(train_x_orig.shape))
# print ("train_y shape: " + str(train_y.shape))
# print ("test_x_orig shape: " + str(test_x_orig.shape))
# print ("test_y shape: " + str(test_y.shape))

train_x_flatten = train_x_orig.reshape(train_x_orig.shape[0], -1).T # (12288, 209)
test_x_flatten = test_x_orig.reshape(test_x_orig.shape[0], -1).T # (12288, 50)

train_x = train_x_flatten / 255.
test_x = test_x_flatten / 255.

print ("train_x's shape: " + str(train_x.shape))
print ("test_x's shape: " + str(test_x.shape))

n_x = 12288
n_h = 7
n_y = 1
layers_dims = (n_x, n_h, n_y)

def two_layer_model(X, Y, layers_dims, learning_rate=0.00075, num_iterations=3000, print_cost=False):
    """

    :param X:
    :param Y:
    :param layers_dims:
    :param learning_rate:
    :param num_iteration:
    :param print_cost:

    :return:
    :parameter -- a dictionary containing W1, b1, W2, b2
    """
    np.random.seed(1)
    grads = {}
    costs = []
    m =X.shape[1]
    (n_x, n_h, n_y) = layers_dims

    parameters = initialize_parameters(n_x, n_h, n_y)

    W1 = parameters["W1"]
    b1 = parameters["b1"]
    W2 = parameters["W2"]
    b2 = parameters["b2"]

    for i in range(num_iterations):
        A1, cache1 = linear_activation_forward(X, W1, b1, "relu")

        A2, cache2 = linear_activation_forward(A1, W2, b2, "sigmoid")


        cost = compute_cost(A2, Y)

        epsilon = 1e-100
        dA2 = -(np.divide(Y, A2+epsilon)) - np.divide(1 - Y, 1 - A2 + epsilon)

        dA1, dW2, db2 = linear_activation_backward(dA2, cache2, activation="sigmoid")
        dA0, dW1, db1 = linear_activation_backward(dA1, cache1, activation="relu")

        grads["dW1"] = dW1
        grads['db1'] = db1
        grads['dW2'] = dW2
        grads['db2'] = db2

        parameters = update_parameters(parameters, grads, learning_rate)

        W1 = parameters["W1"]
        b1 = parameters["b1"]
        W2 = parameters["W2"]
        b2 = parameters["b2"]

        if print_cost and i % 100 == 0:
            print(parameters)
            print("cost after iteration {}: {}".format(i, np.squeeze(cost)))
            costs.append(cost)

    plt.plot(np.squeeze(costs))
    plt.ylabel('cost')
    plt.xlabel('iterations (per tens)')
    plt.title('learning rate = ' + str(learning_rate))
    plt.show()

    return parameters

parameters = two_layer_model(train_x, train_y, layers_dims = (n_x, n_h, n_y), num_iterations = 3000, print_cost=True)
